package com.raos.example.config;

import com.raos.example.func.HighLevelCalculator;
import com.raos.example.service.Assistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * AI助手初始化配置类
 *
 * @author raos
 * @date 2025-03-20
 */
@Configuration
public class AssistantInit {

    @Resource(name = "qwenChatModel")
    ChatLanguageModel chatLanguageModel;

    /**
     * 普通模式AI助手初始化
     */
    /*@Bean(name = "assistant")
    public Assistant init() {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .build();
    }*/

    /**
     * 基于会话记忆的AI助手初始化
     */
    /*@Bean(name = "assistant")
    public Assistant init2() {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .build();
    }*/

    /**
     * 初始化嵌入式存储对象(EmbeddingStore)-临时环境（不建议生产使用，无法共享存储）
     */
    @Bean(name = "inMemoryEmbeddingStore")
    public EmbeddingStore<TextSegment> initEmbeddingStore() {
        // 新建一个基于内存的向量库
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 基于会话隔离的具有RAG能力的AI助手初始化
     */
    @Bean(name = "assistant")
    public Assistant init3(@Autowired @Qualifier("inMemoryEmbeddingStore")
                               EmbeddingStore<TextSegment> embeddingStore) {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();
    }

    @Resource(name = "ollamaChatModel")
    ChatLanguageModel chatLanguageModel2;

    /**
     * 基于会话隔离的具有RAG能力的AI助手初始化
     */
    @Bean(name = "assistant2")
    public Assistant init4(@Autowired @Qualifier("pgVectorEmbeddingStore")
                           EmbeddingStore<TextSegment> embeddingStore) {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel2)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();
    }

    /**
     * 基于会话隔离的具有RAG能力且拥有函数调用能力的AI助手初始化
     */
    @Bean(name = "assistant3")
    public Assistant init5(@Autowired @Qualifier("pgVectorEmbeddingStore")
                           EmbeddingStore<TextSegment> embeddingStore,
                           SearchApiWebSearchEngine searchEngine) {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                // 函数调用能力
                .tools(new HighLevelCalculator(), new WebSearchTool(searchEngine))
                .build();
    }

}
